package com.wg.wgaiagent.app;

import com.wg.wgaiagent.advisor.MyLoggerAdvisor;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.List;

/**
 * @author with葡萄成熟时
 * @version 1.0
 * 普通聊天（包括联网搜索和识别文件）
 */
@Component
@Slf4j
public class ChatApp {


    private final ChatClient chatClient;
    // AI 调用工具能力
    @Resource
    private ToolCallback[] allTools;

    private static final String SYSTEM_PROMPT = "You are a helpful assistant. You will answer any questions users have about the product.";
    /**
     * 初始化 ChatClient
     *
     * @param dashscopeChatModel
     */
    @Autowired
    public ChatApp(ChatModel dashscopeChatModel, InMemoryChatMemoryRepository chatMemoryRepository) {
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder()
                // 使用注入进来的 repository Bean
                .chatMemoryRepository(chatMemoryRepository)
                .maxMessages(20)
                .build();

        chatClient = ChatClient.builder(dashscopeChatModel)
                .defaultSystem(SYSTEM_PROMPT)
                .defaultAdvisors(
                        MessageChatMemoryAdvisor.builder(chatMemory).build(),
                        new MyLoggerAdvisor()
                )
                .build();
    }

    @Resource
    private ToolCallback[] normalChatTools;
    /**
     * 普通聊天功能（支持调用工具）
     *
     * @param message
     * @param chatId
     * @return
     */
    public ChatResponse doChatWithTools(String message, String chatId) {
        ChatResponse chatResponse = chatClient
                .prompt()
                .user(message)
                .advisors(spec -> spec.param(ChatMemory.CONVERSATION_ID, chatId) )
                // 开启日志，便于观察效果
                .advisors(new MyLoggerAdvisor())
                .toolCallbacks(normalChatTools)
                .call()
                .chatResponse();

        return chatResponse;
    }


    /**
     * AI 基础对话（支持多轮对话记忆，SSE 流式传输）
     *
     * @param message
     * @param chatId
     * @return
     */
    public Flux<String> doChatByStream(String message, String chatId) {
        return chatClient
                .prompt()
                .user(message)
                .advisors(spec -> spec.param(ChatMemory.CONVERSATION_ID, chatId))
                .stream()
                .content();
    }

}
